package br.com.config.shard;

import br.com.model.entidade.Customer;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.hibernate.Criteria;
import org.hibernate.Session;
import org.hibernate.SessionFactory;
import org.hibernate.cfg.Configuration;
import org.hibernate.criterion.Restrictions;
import org.hibernate.shards.ShardId;
import org.hibernate.shards.ShardedConfiguration;
import org.hibernate.shards.cfg.ConfigurationToShardConfigurationAdapter;
import org.hibernate.shards.cfg.ShardConfiguration;
import org.hibernate.shards.strategy.ShardStrategy;
import org.hibernate.shards.strategy.ShardStrategyFactory;
import org.hibernate.shards.strategy.ShardStrategyImpl;
import org.hibernate.shards.strategy.access.SequentialShardAccessStrategy;
import org.hibernate.shards.strategy.access.ShardAccessStrategy;
import org.hibernate.shards.strategy.resolution.AllShardsShardResolutionStrategy;
import org.hibernate.shards.strategy.resolution.ShardResolutionStrategy;
import org.hibernate.shards.strategy.selection.ShardSelectionStrategy;

/**
 * The sample application client used for testing
 *
 * @author Arun Viswanathan
 */
public class CustomerApp {

    private SessionFactory sessionFactory;

    public static void main(String[] args) throws Exception {
        CustomerApp app = new CustomerApp();
        app.run();
    }

    private void run() throws SQLException {
        createSchema();
        sessionFactory = createSessionFactory();

        addData();

        Session session = sessionFactory.openSession();
        try {
            Criteria crit = session.createCriteria(Customer.class);
            List count = crit.list();
            System.out.println(count.size());

            crit.add(Restrictions.eq("location", "Bangalore"));
            List reports = crit.list();
            System.out.println(reports.size());

            List customers = getData();
            for (int i = 0; i < customers.size(); i++) {
                Customer customer = (Customer) customers.get(i);
                System.out.println(customer.toString());
            }
        } finally {
            session.close();
        }
    }

    /**
     * Retrieves the data from the Customer tables across the database shards
     *
     * @return result, List of Customer records
     */
    private List getData() {
        Session session = sessionFactory.openSession();
        session.beginTransaction();

        List result = session.createQuery("from Customer").list();
        session.getTransaction().commit();
        session.close();

        return result;
    }

    /**
     * Add Customer details for different Tenants into the database
     */
    private void addData() {
        Session session = sessionFactory.openSession();
        try {
            session.beginTransaction();
            Customer customer = new Customer();
            customer.setFirstName("John");
            customer.setLastName("Doe");
            customer.setLocation("New York City");
            customer.setTenantId("tenant1");
            customer.setUpdateTime(new Date());
            session.save(customer);

            customer = new Customer();
            customer.setFirstName("Jane");
            customer.setLastName("Dane");
            customer.setLocation("San Fransisco");
            customer.setTenantId("tenant2");
            customer.setUpdateTime(new Date());
            session.save(customer);

            customer = new Customer();
            customer.setFirstName("James");
            customer.setLastName("Bond");
            customer.setLocation("London");
            customer.setTenantId("tenant3");
            customer.setUpdateTime(new Date());
            session.save(customer);

            customer = new Customer();
            customer.setFirstName("Vijay");
            customer.setLastName("Raj");
            customer.setLocation("Bangalore");
            customer.setTenantId("tenant4");
            customer.setUpdateTime(new Date());
            session.save(customer);

            customer = new Customer();
            customer.setFirstName("Brett");
            customer.setLastName("Lee");
            customer.setLocation("Melbourne");
            customer.setTenantId("tenant5");
            customer.setUpdateTime(new Date());
            session.save(customer);

            session.getTransaction().commit();
        } finally {
            session.close();
        }
    }

    /**
     * Clean the existing in-memory HSQL database and create new schema with the
     * required tables
     *
     * @throws SQLException
     */
    private void createSchema() throws SQLException {
        for (int i = 1; i < 3; i++) {
            DatabaseUtils.destroyDatabase(i);
            DatabaseUtils.createDatabase(i);
        }
    }

    /**
     * Configure the Hibernate configuration files and create the
     * ShardedSessionFactory
     *
     * @return
     */
    public SessionFactory createSessionFactory() {
        Configuration config = new Configuration();
        config.configure("config/hibernate0.cfg.xml");
        config.addResource("config/customer.hbm.xml");

        // Provide the list of all the database configurations
        List shardConfigs = new ArrayList();
        shardConfigs.add(buildShardConfig("config/hibernate0.cfg.xml"));
        shardConfigs.add(buildShardConfig("config/hibernate1.cfg.xml"));
        shardConfigs.add(buildShardConfig("config/hibernate2.cfg.xml"));

        // Create a virtual shard map of 4 shards which are mapped to 3 physical shards
        Map<Integer, Integer> virtualShardMap = new HashMap<Integer, Integer>();
        virtualShardMap.put(0, 0);
        virtualShardMap.put(1, 1);
        virtualShardMap.put(2, 2);
        virtualShardMap.put(3, 2);

        ShardStrategyFactory shardStrategyFactory = buildShardStrategyFactory();
        ShardedConfiguration shardedConfig = new ShardedConfiguration(
                config,
                shardConfigs,
                shardStrategyFactory,
                virtualShardMap);
        return shardedConfig.buildShardedSessionFactory();
    }

    /**
     * Build the ShardStrategyFactory based on the different shard strategies
     * provided
     *
     * @return
     */
    ShardStrategyFactory buildShardStrategyFactory() {
        return new ShardStrategyFactory() {
            public ShardStrategy newShardStrategy(List<ShardId> shardIds) {
                //RoundRobinShardLoadBalancer loadBalancer = new RoundRobinShardLoadBalancer(shardIds);
                //ShardSelectionStrategy pss = new RoundRobinShardSelectionStrategy(loadBalancer);
                ShardSelectionStrategy pss = new CustomerShardSelectionStrategy();
                ShardResolutionStrategy prs = new AllShardsShardResolutionStrategy(shardIds);
                ShardAccessStrategy pas = new SequentialShardAccessStrategy();
                return new ShardStrategyImpl(pss, prs, pas);
            }
        };
    }

    ShardConfiguration buildShardConfig(String configFile) {
        Configuration config = new Configuration().configure(configFile);
        return new ConfigurationToShardConfigurationAdapter(config);
    }
}
